package yycore

import chisel3._
import chisel3.util._
import chisel3.util.experimental.BoringUtils
import common.Constants._

object CSROpType
{
  // commands
  val CSROpTypeSize = 4
  val N = "b0000".U
  val W = "b0001".U
  val S = "b0010".U
  val C = "b0011".U
  val I = "b0100".U
  val R = "b0101".U

  def apply() = UInt(CSROpTypeSize.W)
}

trait CSRConst {
  // User Trap Setup
  val Ustatus       = 0x000
  val Uie           = 0x004
  val Utvec         = 0x005

  // User Trap Handling
  val Uscratch      = 0x040
  val Uepc          = 0x041
  val Ucause        = 0x042
  val Utval         = 0x043
  val Uip           = 0x044

  // User Floating-Point CSRs (not implemented)
  val Fflags        = 0x001
  val Frm           = 0x002
  val Fcsr          = 0x003

  // User Counter/Timers
  val Cycle         = 0xC00
  val Time          = 0xC01
  val Instret       = 0xC02

  // Supervisor Trap Setup
  val Sstatus_Addr       = 0x100
  val Sedeleg_Addr       = 0x102
  val Sideleg_Addr       = 0x103
  val Sie_Addr           = 0x104
  val Stvec_Addr         = 0x105
  val Scounteren_Addr    = 0x106

  // Supervisor Trap Handling
  val Sscratch_Addr      = 0x140
  val Sepc_Addr          = 0x141.U
  val Scause_Addr        = 0x142
  val Stval_Addr         = 0x143
  val Sip_Addr           = 0x144

  // Supervisor Protection and Translation
  val Satp_Addr          = 0x180.U

  // Machine Information Registers
  val Mvendorid_Addr     = 0xF11.U
  val Marchid_Addr       = 0xF12.U
  val Mimpid_Addr        = 0xF13.U
  val Mhartid_Addr       = 0xF14.U

  // Machine Trap Setup
  val Mstatus_Addr       = 0x300.U
  val Misa_Addr          = 0x301.U
  val Medeleg_Addr       = 0x302.U
  val Mideleg_Addr       = 0x303.U
  val Mie_Addr           = 0x304.U
  val Mtvec_Addr         = 0x305.U
  val Mcounteren_Addr    = 0x306.U

  // Machine Trap Handling
  val Mscratch_Addr      = 0x340.U
  val Mepc_Addr          = 0x341.U
  val Mcause_Addr        = 0x342.U
  val Mtval_Addr         = 0x343.U
  val Mip_Addr           = 0x344.U

  // Machine Memory Protection
  // TBD
  val Pmpcfg0_Addr       = 0x3A0.U
  val Pmpcfg1_Addr       = 0x3A1.U
  val Pmpcfg2_Addr       = 0x3A2.U
  val Pmpcfg3_Addr       = 0x3A3.U
  val PmpaddrBase_Addr   = 0x3B0.U

  val addr_mcycle = 0xb00.U



  def privEcall  = 0x000.U
  def privEbreak = 0x001.U
  def privMret   = 0x302.U
  def privSret   = 0x102.U
  def privUret   = 0x002.U

  def ModeM     = 0x3.U
  def ModeH     = 0x2.U
  def ModeS     = 0x1.U
  def ModeU     = 0x0.U

  def IRQ_UEIP  = 0
  def IRQ_SEIP  = 1
  def IRQ_MEIP  = 3

  def IRQ_UTIP  = 4
  def IRQ_STIP  = 5
  def IRQ_MTIP  = 7

  def IRQ_USIP  = 8
  def IRQ_SSIP  = 9
  def IRQ_MSIP  = 11


  val IntPriority = Seq(
    IRQ_MEIP, IRQ_MSIP, IRQ_MTIP,
    IRQ_SEIP, IRQ_SSIP, IRQ_STIP,
    IRQ_UEIP, IRQ_USIP, IRQ_UTIP
  )
}

trait ExceptionNO {
  def instrAddrMisaligned = 0
  def instrAccessFault    = 1
  def illegalInstr        = 2
  def breakPoint          = 3
  def loadAddrMisaligned  = 4
  def loadAccessFault     = 5
  def storeAddrMisaligned = 6
  def storeAccessFault    = 7
  def ecallU              = 8
  def ecallS              = 9
  def ecallM              = 11
  def instrPageFault      = 12
  def loadPageFault       = 13
  def storePageFault      = 15

  val ExcPriority = Seq(
    breakPoint,
    instrPageFault,
    instrAccessFault,
    illegalInstr,
    instrAddrMisaligned,
    ecallM, ecallS, ecallU,
    storePageFault,
    loadPageFault,
    storeAccessFault,
    loadAccessFault,
    storeAddrMisaligned,
    loadAddrMisaligned
  )
}

object PRV
{
  val SZ = 2
  val U = 0.U
  val S = 1.U
  val H = 2.U
  val M = 3.U
}

class Priv extends Bundle {
  val m = Bool()
  val h = Bool()
  val s = Bool()
  val u = Bool()
}


sealed abstract class CSRBundle extends Bundle with CSRConst with ExceptionNO
sealed abstract class CSRModule extends Module with CSRConst with ExceptionNO

class mstatus extends Bundle{
  val SD = UInt(1.W)
  val reserved = UInt(40.W)
  val TSR = UInt(1.W)
  val TW = UInt(1.W)
  val TVM = UInt(1.W)
  val MXR = UInt(1.W) 
  val SUM = UInt(1.W) 
  val MPRV = UInt(1.W)  
  val XS = UInt(2.W) 
  val FS = UInt(2.W) 
  val MPP = UInt(2.W) 
  val reserved_2 = UInt(2.W) 
  val SPP = UInt(1.W)
  val PIE = new Priv 
  val IE = new Priv 
}
// Machine interrupt-pending register (mip)
class Interrupt extends Bundle {
  val e = new Priv  // external interrupt
  val t = new Priv  // timer interrupt
  val s = new Priv  // software interrupt
}
class mip extends Bundle {
  val reserved = UInt(52.W)
  val intr = new Interrupt
  // val MEIP = UInt(1.W)
  // val reserved_2 = UInt(1.W)
  // val SEIP = UInt(1.W)
  // val UEIP = UInt(1.W)
  // val MTIP = UInt(1.W)
  // val reserved_3 = UInt(1.W)
  // val STIP = UInt(1.W)
  // val UTIP = UInt(1.W)
  // val MSIP = UInt(1.W)
  // val reserved_4 = UInt(1.W)
  // val SSIP = UInt(1.W)
  // val USIP = UInt(1.W)
}
class mie extends Bundle {
  val reserved = UInt(52.W)
  val MEIE = UInt(1.W)
  val reserved_2 = UInt(1.W)
  val SEIE = UInt(1.W)
  val UEIE = UInt(1.W)
  val MTIE = UInt(1.W)
  val reserved_3 = UInt(1.W)
  val STIE = UInt(1.W)
  val UTIE = UInt(1.W)
  val MSIE = UInt(1.W)
  val reserved_4 = UInt(1.W)
  val SSIE = UInt(1.W)
  val USIE = UInt(1.W)
}
class mtvec extends Bundle{
  val BASE = UInt((XLEN - 2).W)
  val MODE = UInt(2.W)
}
class mcause extends Bundle{
  val interrupt = UInt(1.W)
  val exc_code  = UInt((XLEN - 1).W)
}


class CSRReqBundle extends Bundle {
  val pc = Output(UInt(AddrBits.W))
  val cmd = Output(CSROpType())
  val addr = Output(UInt(12.W))
  val wdata = Output(UInt(XLEN.W))
  val rd = Output(UInt(5.W))
  val instValid = Output(Bool())
  val exceptionVec = Output(Vec(16, Bool()))
}

class CSRRespBundle extends Bundle {
  val rdata = Input(UInt(XLEN.W))
  val intrNO = Input(UInt(XLEN.W))

}

class CSRIO extends Bundle {
  val req = Flipped(ValidIO(new CSRReqBundle))
  val resp = ValidIO(new CSRRespBundle)
  val redirect = new RedirectIO
  val wenFix = Output(Bool())
  val raiseIntr = Output(Bool())
}


class CSR() extends CSRModule {
  val io = IO(new CSRIO())


  val write_data = WireInit(0.U(XLEN.W))
  val read_csr_data = WireInit(0.U(XLEN.W))

  val valid = io.req.valid
  val csr_addr = io.req.bits.addr
  val cmd = io.req.bits.cmd
  val pc = io.req.bits.pc
  val instValid = io.req.bits.instValid
  val req_in = io.req.bits
  io.resp.valid := valid


  val mode = RegInit(PRV.M)
  // supervisor
  val reg_satp = RegInit(0.U(XLEN.W))
  val reg_sepc = RegInit(0.U(XLEN.W))

  // Machine-Level CSRs
  // Machine Information Registers
  //val reg_mvendorid,marchid, mimpid, mhartid
  val reg_mhartid   = RegInit(0.U(64.W))
  // Machine Trap Setup
  val reg_mstatus   = RegInit(0.U(XLEN.W))
  val reg_misa = RegInit(0.U(XLEN.W))
  val reg_medeleg = RegInit(0.U(XLEN.W))
  val reg_mideleg = RegInit(0.U(XLEN.W))
  val reg_mie       = RegInit(0.U(XLEN.W))
  val reg_mtvec     = RegInit(0.U.asTypeOf(new mtvec))
  // Machine Trap Handling
  val reg_mscratch  = RegInit(0.U(XLEN.W))
  val reg_mepc      = RegInit(0.U(XLEN.W))
  val reg_mcause    = RegInit(0.U.asTypeOf(new mcause))
  val reg_mtval     = RegInit(0.U(XLEN.W))
  val reg_mip       = RegInit(0.U.asTypeOf(new mip))
  // Machine Memory Protection
  val reg_pmpcfg0 = RegInit(0.U(XLEN.W))
  val reg_pmpaddr0 = RegInit(0.U(XLEN.W))


  val mcycle = RegInit(0.U(64.W))
  mcycle := mcycle + 1.U
  if(!FPGA) {
    BoringUtils.addSource(mcycle, "mcycle")
  }

  val mipWire = WireInit(0.U.asTypeOf(new Interrupt))
  val mip = (mipWire.asUInt | reg_mip.intr.asUInt).asTypeOf(new Interrupt)

  val mstatusStruct = reg_mstatus.asTypeOf(new mstatus)

  // ****************************************
  // READ CSRs
  // Machine Information Registers
  read_csr_data := MuxCase(0.U(XLEN.W), Array(
    (csr_addr === Satp_Addr)        -> reg_satp.asUInt(),
    (csr_addr === Sepc_Addr)        -> reg_sepc.asUInt(),
    (csr_addr === Mhartid_Addr)     -> reg_mhartid.asUInt(),
    (csr_addr === Mstatus_Addr)     -> reg_mstatus,
    (csr_addr === Misa_Addr)        -> reg_misa.asUInt(),
    (csr_addr === Medeleg_Addr)     -> reg_medeleg.asUInt(),
    (csr_addr === Mideleg_Addr)     -> reg_mideleg.asUInt(),
    (csr_addr === Mie_Addr)         -> reg_mie,
    (csr_addr === Mtvec_Addr)       -> reg_mtvec.asUInt(),
    (csr_addr === Mscratch_Addr)    -> reg_mscratch.asUInt(),
    (csr_addr === Mepc_Addr)        -> reg_mepc.asUInt(),
    (csr_addr === Mcause_Addr)      -> reg_mcause.asUInt(),
    (csr_addr === Mtval_Addr)       -> reg_mtval.asUInt(),
    (csr_addr === Mip_Addr)         -> reg_mip.asUInt(),
    (csr_addr === Pmpcfg0_Addr)     -> reg_pmpcfg0.asUInt(),
    (csr_addr === PmpaddrBase_Addr) -> reg_pmpaddr0.asUInt(),
    (csr_addr === addr_mcycle)      -> mcycle
  ))

  io.resp.bits.rdata := Mux(io.req.bits.rd === 0.U, 0.U, read_csr_data)

  write_data := MuxCase(io.req.bits.wdata, Array(
    (cmd === CSROpType.W)     -> io.req.bits.wdata,
    (cmd === CSROpType.S)     -> (read_csr_data | io.req.bits.wdata),
    (cmd === CSROpType.C)     -> (read_csr_data & (~io.req.bits.wdata).asUInt())
  ))

  val system_ins = cmd === CSROpType.I
  val cpu_ren = cmd =/= CSROpType.N && !system_ins
  val wen = valid && cpu_ren && io.req.bits.cmd =/= CSROpType.R
  val opcode = 1.U << csr_addr(2, 0)
  val isEcall = valid && system_ins && opcode(0)
  val is_break = valid && system_ins && opcode(1)
  val insn_ret = valid && system_ins && opcode(2)
  val isMret = valid && insn_ret && !csr_addr(10)

  val raiseExceptionVec = WireInit(io.req.bits.exceptionVec)//Wire(Vec(16, Bool()))
  raiseExceptionVec := DontCare
  raiseExceptionVec(breakPoint) := is_break
  raiseExceptionVec(ecallM) := mode === PRV.M && isEcall
  raiseExceptionVec(ecallU) := mode === PRV.U && isEcall
  val raiseException = raiseExceptionVec.asUInt.orR()
  val exceptionNO = ExcPriority.foldRight(0.U)((i: Int, sum: UInt) => Mux(raiseExceptionVec(i), i.U, sum))

  // Interrupts, 中断优先级：外部中断 > 软件中断 > 计时器中断
  val mtip = WireInit(false.B)
  val msip = WireInit(false.B)
  //val meip = WireInit(false.B)
  BoringUtils.addSink(mtip, "mtip")
  BoringUtils.addSink(msip, "msip")
  //BoringUtils.addSink(meip, "meip")
  mipWire.t.m := mtip
  mipWire.s.m := msip

  val intrVecEnable = Wire(Vec(12, Bool()))
  intrVecEnable.map(_ := ((mstatusStruct.IE.m && mode === PRV.M) || mode < PRV.M))
  val intrVec = reg_mie(11, 0) & mip.asUInt & intrVecEnable.asUInt
  val raiseIntr = intrVec.orR()
  io.raiseIntr := raiseIntr

  // epc for interrupt
//  val if_pc = WireInit(0.U(AddrBits.W))
//  val ibuf_pc = WireInit(0.U(AddrBits.W))
//  val ibuf_instValid = WireInit(false.B)
//  val dec_pc = WireInit(0.U(AddrBits.W))
//  val dec_instValid = WireInit(false.B)
//  BoringUtils.addSink(if_pc, "if_pc")
//  BoringUtils.addSink(ibuf_pc, "ibuf_pc")
//  BoringUtils.addSink(ibuf_instValid, "ibuf_instValid")
//  BoringUtils.addSink(dec_pc, "dec_pc")
//  BoringUtils.addSink(dec_instValid, "dec_instValid")
//  val intrEpc = MuxCase(if_pc, Array(
//    //(io.req.bits.instValidnoComplete === true.B)     -> io.req.bits.pc,
//    (dec_instValid  === true.B)                     -> dec_pc,
//    (ibuf_instValid === true.B)                     -> ibuf_pc
//  ))
  val intrEpc = io.req.bits.pc
  val epcNew = Mux(raiseIntr, intrEpc, io.req.bits.pc)
  val intrNO = IntPriority.foldRight(0.U)((i: Int, sum: UInt) => Mux(intrVec(i), i.U, sum))
  val causeNO = (raiseIntr << (XLEN-1)).asUInt() | Mux(raiseIntr, intrNO, exceptionNO)

  val raiseExceptionIntr = (raiseException || raiseIntr) && instValid
  when(raiseExceptionIntr) {
    val old_mstatus = WireInit(reg_mstatus.asTypeOf(new mstatus))
    val new_mstatus = WireInit(reg_mstatus.asTypeOf(new mstatus))
    reg_mcause := causeNO.asTypeOf(new mcause)
    reg_mepc := epcNew
    new_mstatus.MPP := mode
    new_mstatus.PIE.m := old_mstatus.IE.m
    new_mstatus.IE.m := false.B
    mode := PRV.M

    reg_mstatus := new_mstatus.asUInt()
  }
  io.resp.bits.intrNO := Mux(raiseIntr, causeNO, 0.U)
  io.wenFix := raiseException

  val retTarget = Wire(UInt(AddrBits.W))
  retTarget := DontCare
  // mret
  when(isMret){
    val old_mstatus = WireInit(reg_mstatus.asTypeOf(new mstatus))
    val new_mstatus = WireInit(reg_mstatus.asTypeOf(new mstatus))
    new_mstatus.IE.m := old_mstatus.PIE.m
    new_mstatus.PIE.m := true.B
    new_mstatus.MPP := PRV.U
    new_mstatus.MPRV := 0.U
    mode := old_mstatus.MPP
    // update reg status
    reg_mstatus := new_mstatus.asUInt()
    retTarget := reg_mepc
  }

  val isXret = isMret

  val trapTarget = Mux(reg_mtvec.MODE === 0.U, Cat(reg_mtvec.BASE, Fill(2, 0.U)),
    Cat(reg_mtvec.BASE, Fill(2, 0.U)) + reg_mcause.asUInt() << 2.U)//reg_mtvec
  io.redirect.valid := raiseExceptionIntr || isXret
  io.redirect.target := Mux(raiseExceptionIntr, trapTarget, retTarget)

  // write data
  when(wen)
  {
    when(csr_addr === Satp_Addr){
      reg_satp := write_data
    }.elsewhen(csr_addr === Sepc_Addr){
      reg_sepc := write_data
    }.elsewhen(csr_addr === Mhartid_Addr){
      reg_mhartid := write_data
    }.elsewhen(csr_addr === Mstatus_Addr){
      val wMstatus = WireInit(write_data.asTypeOf(new mstatus))
      reg_mstatus := Cat(wMstatus.FS === "b11".U, write_data(XLEN-2, 0))
    }.elsewhen(csr_addr === Misa_Addr){
      reg_misa := write_data
    }.elsewhen(csr_addr === Medeleg_Addr){
      reg_medeleg := write_data
    }.elsewhen(csr_addr === Mideleg_Addr){
      reg_mideleg := write_data
    }.elsewhen(csr_addr === Mie_Addr){
      reg_mie := write_data
    }.elsewhen(csr_addr === Mtvec_Addr){
      reg_mtvec := write_data.asTypeOf(new mtvec)
    }.elsewhen(csr_addr === Mscratch_Addr) {
      reg_mscratch := write_data
    }.elsewhen(csr_addr === Mepc_Addr){
      reg_mepc := write_data
    }.elsewhen(csr_addr === Mcause_Addr){
      reg_mcause := write_data.asTypeOf(new mcause)
    }.elsewhen(csr_addr === Mtval_Addr){
      reg_mtval := write_data
    }.elsewhen(csr_addr === Mip_Addr){
      reg_mip := write_data.asTypeOf(new mip)
    }.elsewhen(csr_addr === Mepc_Addr){
      reg_mepc := write_data
    }.elsewhen(csr_addr === Pmpcfg0_Addr){
      reg_pmpcfg0 := write_data
    }.elsewhen(csr_addr === PmpaddrBase_Addr){
      reg_pmpaddr0 := write_data
    }

    // for differential testing
    if(!FPGA){
//      BoringUtils.addSource(raiseIntr, "difftestIntr")
      BoringUtils.addSource(RegNext(mode), "difftestMode")
      BoringUtils.addSource(RegNext(reg_mstatus), "difftestMstatus")
      //  BoringUtils.addSource(RegNext(reg_sstatus), "difftestSstatus")
      BoringUtils.addSource(RegNext(reg_mepc.asUInt()), "difftestMepc")
      //BoringUtils.addSource(RegNext(reg_sepc), "difftestSepc")
      BoringUtils.addSource(RegNext(reg_mcause.asUInt()), "difftestMcause")
    }

  }
}
